﻿

#include "ArmorDetector.h"

//#define Debug  //调试使用
//#define Calculate_time //调试程序用时
float times=0.4;
ArmorDetector::ArmorDetector()
{
    now_frame=0;
    num_infantry.resize(0);
    direction=0;

    enable_debug = true;
    //number_template.emplace_back(cv::Mat());
    for (int i = 0; i <= 8; ++i)
        number_template.emplace_back(cv::imread(std::to_string(i) + ".png", cv::IMREAD_GRAYSCALE));
    
    //1.kalman filter setup
    KF.init(stateNum, measureNum, 0);
    KF.transitionMatrix = (cv::Mat_<float>(4, 4) << 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1);  //转移矩阵A
    setIdentity(KF.measurementMatrix);                                             //测量矩阵H
    setIdentity(KF.processNoiseCov, cv::Scalar::all(1e-5));                            //系统噪声方差矩阵Q
    setIdentity(KF.measurementNoiseCov, cv::Scalar::all(1e-1));                        //测量噪声方差矩阵R
    setIdentity(KF.errorCovPost, cv::Scalar::all(1));                                  //后验错误估计协方差矩阵P
    rng.fill(KF.statePost, cv::RNG::UNIFORM, 0, 480);   //初始状态值x(0)
    measurement = cv::Mat::zeros(measureNum, 1, CV_32F);                           //初始测量值，因为后面要更新这个值，所以必须先定义  
    svm = cv::ml::SVM::create();
    svm = cv::Algorithm::load<cv::ml::SVM>("svm.xml");
}

ArmorDetector::~ArmorDetector()
{
}

bool ArmorDetector::DetectArmor(cv::Mat& src, cv::Point3f& target)
{
#ifdef Calculate_time

    static int times = 0;
    static double sum = 0;
    double t = (double)cv::getTickCount();//计算程序耗时

#endif

#ifdef Debug

    if (enable_debug) {
        show_lights = src.clone();
        show_armors = src.clone();
        show_final_armor = src.clone();
        //printf("%d\n", mcu_data.enemy_color);
        cv::imshow("src", src);
        auto c = cv::waitKey(1);
        if (c == 'a') {
            cv::waitKey();
        }
    }

#endif

    std::vector<cv::RotatedRect> lights;
    DetectLights(src, lights);/*灯条判定*/

    std::vector<std::vector<cv::Point2f>> armors;
    PossibleArmors(lights, armors);/*装甲板判定*/
    if (armors.empty()) {
        return false;
    }

    std::vector<cv::Point2f> final_armor;
    int num;
    SelectFinalArmor(src, armors, final_armor,num);/*装甲板选择*/
    if (final_armor.empty()) {
        return false;
    }
//    std::cout<<num<<std::endl;
    static double left_ave=0.0, right_ave=0.0;
    if(true)
    {
        if(direction==0&&num_infantry.size()<=15)
        {
            num_infantry.push_back((final_armor[0].x+final_armor[1].x)/2);
            if(num_infantry.size()<15) return false;
            else
            {
                int score1,score2;//left and right
          
                score1=score2=0;
                for(int i=5;i<num_infantry.size()-1;i++)
                {
                    if(num_infantry[i]<=num_infantry[i+1]){
                       score1++;
                       right_ave=right_ave+(num_infantry[i+1]-num_infantry[i]);
                    }
                    else {
                       score2++;
                       left_ave=left_ave+(num_infantry[i]-num_infantry[i+1]);
                    }
                 printf("left:%lf right:%lf\n",left_ave,right_ave);
                }
                right_ave/=4.0;
                left_ave/=4.0;
                if(score1>=score2) direction=1;
                else direction=2;
                num_infantry.clear();
                num_infantry.resize(0);
                return false;
            }
        }
        else
        {
            if(now_frame>130)
            {
                direction=0;
                now_frame=0;
                left_ave=0.0;
                right_ave=0.0;
                return false;
            }
            else
            {  // printf("right:%lf left:%lf\n",right_ave,left_ave);
                now_frame++;
                float width=std::abs(final_armor[0].x-final_armor[1].x);
                if(direction==2)
                {
                    final_armor[0].x-=(final_armor[0].x-width*times*left_ave)<0?final_armor[0].x:width*times*left_ave;
                    final_armor[3].x-=(final_armor[3].x-width*times*left_ave)<0?final_armor[3].x:width*times*left_ave;
                    final_armor[1].x=final_armor[0].x+width;
                    final_armor[2].x=final_armor[3].x+width;
                }
                else if(direction==1)
                {
                    final_armor[1].x=(final_armor[1].x+width*times*right_ave)>cameraParam.resolution_width?cameraParam.resolution_width:(width*times*right_ave+final_armor[1].x);
                    final_armor[2].x=(final_armor[2].x+width*times*right_ave)>cameraParam.resolution_width?cameraParam.resolution_width:(width*times*right_ave+final_armor[2].x);
                    final_armor[0].x=final_armor[2].x-width;
                    final_armor[3].x=final_armor[1].x-width;
                }
                else std::cout<<"direction error!"<<std::endl;
                CalcControlInfo(final_armor, target);
                
                for (int i = 0; i < 4; i++)
                cv::line(src, final_armor[i], final_armor[(i + 1) % 4], cv::Scalar(0, 255, 0), 2, cv::LINE_AA);
                cv::imshow("src",src);
                cv::waitKey(10);
                return true;
            }
        }
    }
    else ;

    //CalsolvePnP(final_armor, target);
    CalcControlInfo(final_armor, target);/*计算*/
    //Kalman_Predict(final_armor, target);//预测

#ifdef Calculate_time

    t = (double)cv::getTickCount() - t;
    sum += t * 1000. / cv::getTickFrequency();
    times++;
    if (times == 1000)
    {
        printf("execution time = %gms\n", sum / times);//转换时间单位并输出代码运行时间
        times = 0;
        sum = 0;
    }

#endif
        
    return true;
}

cv::Mat ArmorDetector::DistillationColor(const cv::Mat& src)/*颜色二值化*/
{
    std::vector<cv::Mat> bgr;
    cv::split(src, bgr);
    cv::Mat result;
    if (mcu_data.enemy_color == EnemyColor::BLUE)
        cv::subtract(bgr[0], bgr[2], result);
    else
        cv::subtract(bgr[2], bgr[0], result);
    return result;
}

inline void ArmorDetector::DrawRotatedRect(const cv::Mat& image, const cv::RotatedRect& rect, const cv::Scalar& color, int thickness)/*标记矩形函数*/
{
    cv::Point2f vertices[4];
    rect.points(vertices);
    for (int i = 0; i < 4; i++)
        cv::line(image, vertices[i], vertices[(i + 1) % 4], color, thickness, cv::LINE_AA);
}

void ArmorDetector::DetectLights(const cv::Mat& src, std::vector<cv::RotatedRect>& lights)
{
    cv::Mat gray_image;
    cv::cvtColor(src, gray_image, cv::COLOR_BGR2GRAY);
    cv::Mat binary_brightness_image;
    float thresh_brightness = (mcu_data.enemy_color == EnemyColor::BLUE) ? armorParam.brightness_thresh_blue : armorParam.brightness_thresh_red;
    cv::threshold(gray_image, binary_brightness_image, thresh_brightness, 255, cv::THRESH_BINARY);

    cv::Mat binary_color_image;
    cv::Mat distillation = DistillationColor(src);

    float thresh_color = (mcu_data.enemy_color == EnemyColor::BLUE) ? armorParam.blue_thresh : armorParam.red_thresh;
    cv::threshold(distillation, binary_color_image, thresh_color, 255, cv::THRESH_BINARY);
    // cv::morphologyEx(binary_color_image, binary_color_image, cv::MORPH_CLOSE, cv::Mat());
    cv::dilate(binary_color_image, binary_color_image, cv::Mat());  /*膨胀*/

    std::vector<std::vector<cv::Point>> contours_color; /*颜色轮廓*/
    cv::findContours(binary_color_image, contours_color, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
    std::vector<std::vector<cv::Point>> contours_brightness; /*亮度轮廓*/
    cv::findContours(binary_brightness_image, contours_brightness, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    lights.reserve(contours_color.size());
    for (unsigned int i = 0; i < contours_brightness.size(); ++i) {
        for (unsigned int j = 0; j < contours_color.size(); ++j) {
            if (cv::pointPolygonTest(contours_color[j], contours_brightness[i][0], false) >= 0.0) {     /*判断亮度点是否在颜色轮廓内*/
                cv::RotatedRect single_light = cv::minAreaRect(contours_brightness[i]);
                cv::Rect bounding_rect = single_light.boundingRect();/*得到一个垂直的矩形*/
                if (bounding_rect.height < bounding_rect.width || single_light.size.area() < armorParam.light_min_area)    /*垂直矩形的高度大于宽度，且面积大于最小值*/
                    break;
                lights.emplace_back(single_light);
                break;
            }
        }
    }

#ifdef Debug

    if (enable_debug) {
        for (auto light : lights)
            DrawRotatedRect(show_lights, light, cv::Scalar(0, 255, 0), 2);
        cv::imshow("distillation", distillation);
        cv::imshow("binary_brightness_image", binary_brightness_image);
        cv::imshow("binary_color_img", binary_color_image);
        cv::imshow("show_lights", show_lights);
    }

#endif

}

inline void ArmorDetector::DrawArmor(const cv::Mat& image, const std::vector<cv::Point2f>& armor, const cv::Scalar& color, int thickness)/*标注装甲板*/
{
    for (int i = 0; i < 4; i++)
        cv::line(image, armor[i], armor[(i + 1) % 4], color, thickness, cv::LINE_AA);
}

void ArmorDetector::PossibleArmors(const std::vector<cv::RotatedRect>& lights, std::vector<std::vector<cv::Point2f>>& armors)
{
    for (unsigned int i = 0; i < lights.size(); ++i) {
        for (unsigned int j = i + 1; j < lights.size(); ++j) {
            float angle_i = (lights[i].angle <= -45.0) ? lights[i].angle + 90 : lights[i].angle;
            float angle_j = (lights[j].angle <= -45.0) ? lights[j].angle + 90 : lights[j].angle;
            float angle_diff = std::abs(angle_i - angle_j);                   /*由于opencv的角度选择随矩形旋转角度不同并不一致，为方便处理这里统一进行变换*/
             //std::cout << "\nangle_diff = " << angle_diff;
            if (angle_diff > armorParam.armor_max_angle_diff)
                continue;

            double armor_angle = (double)std::atan((lights[i].center.y - lights[j].center.y) / (lights[i].center.x - lights[j].center.x)) * 180 / CV_PI;
            // std::cout << "\nangle_i = " << angle_i;
            // std::cout << "\nangle_j = " << angle_j;
            // std::cout << "\narmor_angle = " << armor_angle;
            if (std::max(abs(armor_angle - angle_i), abs(armor_angle - angle_j)) > armorParam.armor_max_angle_diff) /*角度过大*/
                continue;

            float height_i = std::max(lights[i].size.width, lights[i].size.height);
            float height_j = std::max(lights[j].size.width, lights[j].size.height);
            auto height = std::minmax(height_i, height_j);
            //std::cout << "\nheight_i/height_j = " << height.second / height.first << '\n';
            if ((height.second / height.first) > armorParam.armor_max_height_ratio)  /*两个灯条的高度差*/
                continue;

            float armor_width = std::sqrt((lights[i].center.x - lights[j].center.x) * (lights[i].center.x - lights[j].center.x) + (lights[i].center.y - lights[j].center.y) * (lights[i].center.y - lights[j].center.y));
            float armor_height = (height_i + height_j) / 2;
            // std::cout << "\narmor_width = " << narmor_width;
            // std::cout << "\narmor_height = " << armor_height;
           // std::cout << "\nwidth/height = " << armor_width / armor_height;
            if ((armor_width / armor_height) > armorParam.armor_max_aspect_ratio)  /*装甲板的宽高比*/
                continue;

            std::vector<cv::Point2f> armor;
            if (lights[i].center.x < lights[j].center.x)
                CalcArmorInfo(armor, lights[i], lights[j]);
            else
                CalcArmorInfo(armor, lights[j], lights[i]);  
            armors.emplace_back(armor);
        }
    }

#ifdef Debug

    if (enable_debug) {
        for (auto armor : armors)
            DrawArmor(show_armors, armor, cv::Scalar(0, 255, 0), 2);
        cv::imshow("show_armors", show_armors);
    }

#endif

}

void ArmorDetector::CalcArmorInfo(std::vector<cv::Point2f>& armor, const cv::RotatedRect& left_light, const cv::RotatedRect& right_light) 
/*计算装甲板的角点*/
{
    cv::Point2f left_points[4], right_points[4];
    left_light.points(left_points);
    right_light.points(right_points);

    std::sort(left_points, left_points + 4, [](const cv::Point2f& p1, const cv::Point2f& p2) { return p1.y < p2.y; });
    cv::Point2f left_up = (left_points[0].x > left_points[1].x) ? left_points[0] : left_points[1];
    cv::Point2f left_down = (left_points[2].x > left_points[3].x) ? left_points[2] : left_points[3];
    cv::Point2f left_shift = (left_down - left_up) * armorParam.shift_scale;

    std::sort(right_points, right_points + 4, [](const cv::Point2f& p1, const cv::Point2f& p2) { return p1.y < p2.y; });
    cv::Point2f right_up = (right_points[0].x < right_points[1].x) ? right_points[0] : right_points[1];
    cv::Point2f right_down = (right_points[2].x < right_points[3].x) ? right_points[2] : right_points[3];
    cv::Point2f right_shift = (right_down - right_up) * armorParam.shift_scale;

    armor.emplace_back(left_up - left_shift);
    armor.emplace_back(right_up - right_shift);
    armor.emplace_back(right_down + right_shift);
    armor.emplace_back(left_down + left_shift);
}

int ArmorDetector::DetectNumber(const cv::Mat& perspective) /*判断数字，方法为SVM*/
{
    cv::Mat tmp, tmp1;
    cv::resize(perspective, tmp, cv::Size(10, 10));
    tmp.convertTo(tmp1, CV_32FC1);
    float result = svm->predict(tmp1.reshape(1, 1));
    return (int)result;
}

void ArmorDetector::SelectFinalArmor(const cv::Mat& src, std::vector<std::vector<cv::Point2f>>& armors, std::vector<cv::Point2f>& final_armor,int &num)
{
    //判断标签
    cv::Mat gray;
    cv::cvtColor(src, gray, cv::COLOR_BGR2GRAY);
    cv::threshold(gray, gray, armorParam.number_thresh, 255, cv::THRESH_BINARY);

    static float half_camera_width = cameraParam.resolution_width / 2;
    static float half_camera_height = cameraParam.resolution_height / 2;
    std::vector<float> armors_distance(armors.size()); //armors_distance 距离中心的距离
    for (int i = 0; i < armors.size(); i++) {
        cv::Mat perspective;
        cv::Mat homography = cv::findHomography(armors[i], armor_points, cv::RANSAC);
        cv::warpPerspective(gray, perspective, homography, cv::Size(armorParam.small_armor_width, armorParam.armor_height));/*该处为矩阵变换*/

        int n = DetectNumber(perspective);
        if (n != 0) {
            cv::Point2f armor_center = (armors[i][0] + armors[i][1] + armors[i][2] + armors[i][3]) / 4;
            armors_distance[i] = sqrt(powf(armor_center.x - half_camera_width, 2) + powf(armor_center.y - half_camera_height, 2));
            
#ifdef Debug
            
           if (enable_debug) {
                printf("%d\n", n);
                cv::imshow("SVM", perspective);
        }

#endif
        }
        else
            armors_distance[i] = 8;
        
    }
    int max_level_index = std::distance(armors_distance.begin(), std::min_element(armors_distance.begin(), armors_distance.end()));/*选择距离中心最近的目标*/
    
    if (armors_distance[max_level_index] == 800) //表明只有错误装甲板，返回空
        return;

    final_armor = armors[max_level_index];
    
    cv::Mat perspective;
    cv::Mat homography = cv::findHomography(armors[max_level_index], armor_points, cv::RANSAC);
    cv::warpPerspective(gray, perspective, homography, cv::Size(armorParam.small_armor_width, armorParam.armor_height));/*该处为矩阵变换*/
    num = DetectNumber(perspective);

#ifdef Debug

    if (enable_debug) {
        DrawArmor(show_final_armor, armors[max_level_index], cv::Scalar(0, 255, 0), 2);
        cv::imshow("show_final_armor", show_final_armor);
    }

#endif

}

void ArmorDetector::CalcControlInfo(const std::vector<cv::Point2f>& armor, cv::Point3f& target)/*相似三角形计算角度，距离*/
{
    static float half_camera_width = cameraParam.resolution_width / 2;
    static float half_camera_height = cameraParam.resolution_height / 2;
    static float camera_fx = cameraParam.fx;
    static float camera_fy = cameraParam.fy;
    constexpr double rad = 57.295779513082320876798154814105;

    cv::Point2f armor_up = (armor[0] + armor[1]) / 2;
    cv::Point2f armor_down = (armor[2] + armor[3]) / 2;
    float armor_height = sqrt(powf(armor_down.x - armor_up.x, 2) + powf(armor_down.y - armor_up.y, 2));
    cv::Point2f armor_center = (armor_up + armor_down) / 2;

    float yaw_offset = atan2(half_camera_width - armor_center.x, camera_fx) * rad;
    float pitch_offset = atan2(half_camera_height - armor_center.y, camera_fy) * rad;
    float distance = camera_fx * armorParam.armor_height / armor_height / 1000;
    //std::cout << "\nyaw=" << yaw_offset;
    //std::cout << "\npitch=" << pitch_offset;
    //std::cout << "\ndistance" << distance;
    target = cv::Point3f(yaw_offset, pitch_offset, distance);
}

void ArmorDetector::CalsolvePnP(const std::vector<cv::Point2f>& armor, cv::Point3f& target)
{   
    std::vector<cv::Point3f> objP;
    cv::Mat objM;
    objP.clear();
    objP.push_back(cv::Point3f(0, 0, 0));
    objP.push_back(cv::Point3f(150, 0, 0));
    objP.push_back(cv::Point3f(150, 150, 0));
    objP.push_back(cv::Point3f(0, 150, 0));
    cv::Mat(objP).convertTo(objM, CV_32F);
    cv::Mat rvec, tvec;
    cv::solvePnP(objM, cv::Mat(armor), NULL, NULL, rvec, tvec, false, 0);
}

void ArmorDetector::Kalman_Predict(const std::vector<cv::Point2f>& armor, cv::Point3f& target)
{
    float sample_time = 10.0;
    //2.kalman prediction
    cv::Mat prediction = KF.predict();
    cv::Point2f predict_pt = cv::Point2f(prediction.at<float>(0) + prediction.at<float>(2) * sample_time, prediction.at<float>(1) + prediction.at<float>(3) * sample_time);   //预测值(x',y')


    cv::Point2f armor_center = (armor[0] + armor[1] + armor[2] + armor[3]) / 4;
    //3.update measurement
    measurement.at<float>(0) = (float)armor_center.x;
    measurement.at<float>(1) = (float)armor_center.y;

    //4.update
    KF.correct(measurement);	

    static float half_camera_width = cameraParam.resolution_width / 2;
    static float half_camera_height = cameraParam.resolution_height / 2;
    static float camera_fx = cameraParam.fx;
    static float camera_fy = cameraParam.fy;
    constexpr double rad = 57.295779513082320876798154814105;

    float yaw_offset = atan2(half_camera_width - predict_pt.x, camera_fx) * rad;
    float pitch_offset = atan2(half_camera_height - predict_pt.y, camera_fy) * rad;

    target.x = yaw_offset;
    target.y = pitch_offset;

#ifdef Debug

    if (enable_debug){

        //draw 
        cv::circle(show_armors, predict_pt, 5, cv::Scalar(0, 0, 255), 3);    //predicted point with green
        cv::circle(show_armors, armor_center, 5, cv::Scalar(255, 0, 0), 3); //current position with red	
       
        char buf[256];
#ifdef Windows
        sprintf_s(buf, sizeof(buf), "predicted position:(%3f,%3f)", predict_pt.x, predict_pt.y);
        cv::putText(show_armors, buf, cv::Point(10, 30), CV_FONT_HERSHEY_SCRIPT_COMPLEX, 1, cv::Scalar(0, 0, 0), 1, 8);
        sprintf_s(buf, sizeof(buf), "predicted speed:(%3f,%3f)", prediction.at<float>(2), prediction.at<float>(3));
        cv::putText(show_armors, buf, cv::Point(10, 60), CV_FONT_HERSHEY_SCRIPT_COMPLEX, 1, cv::Scalar(0, 0, 0), 1, 8);
        sprintf_s(buf, sizeof(buf), "current position :(%3f,%3f)", armor_center.x, armor_center.y);
        cv::putText(show_armors, buf, cv::Point(10, 90), CV_FONT_HERSHEY_SCRIPT_COMPLEX, 1, cv::Scalar(0, 0, 0), 1, 8);

#elif defined Linux
        snprintf(buf, sizeof(buf), "predicted position:(%3f,%3f)", predict_pt.x, predict_pt.y);
        cv::putText(show_armors, buf, cv::Point(10, 30), CV_FONT_HERSHEY_SCRIPT_COMPLEX, 1, cv::Scalar(0, 0, 0), 1, 8);
        snprintf(buf, sizeof(buf), "predicted speed:(%3f,%3f)", prediction.at<float>(2), prediction.at<float>(3));
        cv::putText(show_armors, buf, cv::Point(10, 60), CV_FONT_HERSHEY_SCRIPT_COMPLEX, 1, cv::Scalar(0, 0, 0), 1, 8);
        snprintf(buf, sizeof(buf), "current position :(%3f,%3f)", armor_center.x, armor_center.y);
        cv::putText(show_armors, buf, cv::Point(10, 90), CV_FONT_HERSHEY_SCRIPT_COMPLEX, 1, cv::Scalar(0, 0, 0), 1, 8);
#endif
       
        cv::imshow("Kalman", show_armors);
    }
#endif
   

}
